import numpy as np
import cv2
def filter_mask(mask, threshold=0.5):
    """
    过滤掩码中小于阈值的区域，生成二值化掩码。
    
    参数：
        mask (numpy.ndarray): 输入的注意力掩码，形状为 (H, W)，值通常在 [0, 1] 区间
        threshold (float): 阈值，小于此值的区域置为 0，大于等于此值的区域置为 1
    
    返回：
        binary_mask (numpy.ndarray): 二值化掩码，形状为 (H, W)，值在 {0, 1}
    """
    # 确保 mask 是浮点类型且在 [0, 1] 范围内
    mask = np.clip(mask, 0, 1)
    # 应用阈值，二值化
    binary_mask = (mask >= threshold).astype(np.uint8)
    return binary_mask


def contour(binary_mask):
    """
    对二值化掩码进行轮廓分析，提取完整区域的边缘。
    
    参数：
        binary_mask (numpy.ndarray): 二值化掩码，形状为 (H, W)，值在 {0, 1}
    
    返回：
        contours (list): 轮廓列表，每个轮廓是一个形状为 (N, 1, 2) 的数组，表示边缘点坐标
    """
    # 查找轮廓
    contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    return contours
def visual_contour(image_path, contours, save_path):
    """
    使用给定的轮廓信息，在原始图像上绘制轮廓框并保存。
    
    参数：
        image_path (str): 原始图像的文件路径
        contours (list): 轮廓列表，由 contour 函数返回
        save_path (str): 保存绘制结果的文件路径
    """
    # 读取原始图像
    img = cv2.imread(image_path)
    if img is None:
        raise ValueError(f"无法读取图像：{image_path}")
    
    # 绘制轮廓
    cv2.drawContours(img, contours, -1, (0, 255, 0), 2)  # 绿色线条，厚度为 2
    
    # 保存结果
    cv2.imwrite(save_path, img)
    print(f"已保存可视化结果到：{save_path}")
    
if __name__ == '__main__':
    import os
    mask_img_dir = './experiments/test/data/mask'
    src_img_dir = '/mnt/c/DocumentWorkSpace/public_processes/APTOS/images'
    tar_visual_dir = './experiments/test/data/contour'
    os.makedirs(tar_visual_dir, exist_ok=True)
    
    for mask_name in os.listdir(mask_img_dir):
        mask_path = os.path.join(mask_img_dir, mask_name)
        src_img_path = os.path.join(src_img_dir, mask_name[7:])
        tar_visual_path = os.path.join(tar_visual_dir, mask_name)
        
        # 读取灰度掩码，值范围为 [0, 255]
        mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
        if mask is None:
            print(f"无法读取掩码：{mask_path}")
            continue
        
        # 归一化到 [0, 1]
        mask = mask / 255.0
        
        # 应用阈值并提取轮廓
        binary_mask = filter_mask(mask, threshold=0.5)
        contours = contour(binary_mask)
        visual_contour(src_img_path, contours, tar_visual_path)